package shuanfa.tree;

import com.sun.org.apache.regexp.internal.RE;

public class RBTree<K extends Comparable<K>, V> {
    private static final boolean RED = false;
    private static final boolean BLOCK = true;

    private RBNode root;

    public RBNode getRoot() {
        return root;
    }

    /**
     * 红黑树新增节点
     * 1.普通二叉树的插入
     * 查询到插入的位置 将k，v对封装成RBTree节点 插入树中
     * 2.红黑树的平衡（调整  旋转+变色）
     */
    public void put(K key, V value) {
        if (key == null) throw new NullPointerException();
        RBNode t = this.root;
        if (t == null) {
            t = new RBNode<>(key, value == null ? key : value, null);
            setColor(t, BLOCK);
            return;
        }
        int cmp;//定义插入的位置
        RBNode parent;//定义寻找节点的父节点
        do {
            parent = t;
            cmp = key.compareTo((K) t.key);
            if (cmp < 0) {
                t = t.left;
            } else if (cmp > 0) {
                t = t.right;
            } else {
                t.setValue(value == null ? key : value);
            }
        }
        while (t != null);
        //找到插入位置 parent的子节点
        RBNode<K, Object> rbNode = new RBNode<>(key, value == null ? key : value, null);
        if (cmp < 0) {
            parent.left = rbNode;
        } else {
            parent.right = rbNode;
        }
        //做平衡处理，旋转，变色
        fixAfterPut(rbNode);

    }

    public V remove(K key){
        //根据key获取节点
        RBNode node = getNode(key);
        if (node ==null){
            return null;
        }
        //把删除节点原来的值保存下来
        V oldValue = (V)node.value;
        //删除节点的方法
        deleteNode(node);
        return oldValue;
    }

    //删除节点
    private void deleteNode(RBNode node) {
        if (node == null) return;
        //如果删除的节点有两个子节点，用前驱或后继节点替代
        if (leftOf(node) != null && rightOf(node)!=null){
            //找到要删除的节点的前驱节点
            RBNode predecessor = predecessor(node);
            node.key = predecessor.key;
            node.value = predecessor.value;
            //要删除的节点 前驱节点
            node = predecessor;
        }
        //如果节点有一个子节点，用子节点替代
        else if (leftOf(node) == null && rightOf(node) != null){
            node.key = rightOf(node).key;
            node.value = rightOf(node).value;
            node = rightOf(node);
        }else if (leftOf(node) != null && rightOf(node) == null){
            node.key = leftOf(node).key;
            node.value = leftOf(node).value;
            node = leftOf(node);
        }
        RBNode replacement = node;
        //如果是叶子节点 直接删除
        //删除匹配到的node节点
        //等待GC回收node
        node.left = node.right =node.parent =null;
        node = null;
        //删除后平衡变色
        if (replacement.color == BLOCK){
            fixAfterRemove(replacement);
        }



    }
    //删除后平衡变色
    private void fixAfterRemove(RBNode replacement) {

    }

    //根据key获取节点
    private RBNode getNode(K key) {
        if (key == null) throw new NullPointerException();
        RBNode node = this.root;
        int cmp = 0;
        while (node != null){
            cmp = key.compareTo((K) root.key);
            if (cmp>0){
                node = node.right;
            }else if (cmp<0){
                node = node.left;
            }else {
                return node;
            }
        }
        return null;
    }

    /**
     * 寻找前驱节点
     *
     * @param rbNode
     * @return
     */
    private RBNode predecessor(RBNode rbNode) {
        if (rbNode == null) return null;
        RBNode right = null;
        RBNode left = rbNode.left;
        do {
            if (left != null) {
                right = left.right;
            }else {
                return null;
            }
        } while (right != null);
        return right.parent;
    }

    private boolean colorOf(RBNode rbNode) {
        return rbNode == null ? BLOCK : rbNode.color;
    }

    private RBNode parentOf(RBNode rbNode) {
        return rbNode != null ? rbNode.parent : null;
    }

    private RBNode leftOf(RBNode rbNode) {
        return rbNode != null ? rbNode.left : null;
    }

    private RBNode rightOf(RBNode rbNode) {
        return rbNode != null ? rbNode.right : null;
    }

    private void setColor(RBNode rbNode, boolean color) {
        if (rbNode != null) {
            rbNode.setColor(color);
        }
    }

    /**
     * 红黑树平衡处理，旋转，变色
     *
     * @param rbNode
     */
    private void fixAfterPut(RBNode rbNode) {
        //插入的节点为红色节点
        setColor(rbNode, RED);
        //8种情况需要处理 4种处理方式
        //循环递归处理
        while (rbNode != null && rbNode != root && rbNode.color == RED) {
            //插入的节点的父节点 为 父节点的父节点的左节点
            if (parentOf(rbNode) == parentOf(parentOf(rbNode)).left) {
                RBNode y = rightOf(parentOf(parentOf(rbNode)));//叔叔节点
                if (colorOf(y) == RED) {//如果叔叔节点存在，且为红
                    //变色处理
                    setColor(parentOf(rbNode), BLOCK);
                    setColor(y, BLOCK);
                    setColor(parentOf(y), RED);
                    //递归处理
                    rbNode = parentOf(parentOf(rbNode));
                } else {
                    //如果呈现三角形
                    if (rightOf(parentOf(rbNode)) == rbNode) {
                        //对父节点左旋
                        leftRBNode(parentOf(rbNode));
                    }
                    //对爷爷节点 右旋  变色
                    setColor(parentOf(rbNode), BLOCK);
                    setColor(parentOf(parentOf(rbNode)), RED);
                    rightRBNode(parentOf(parentOf(rbNode)));
                }
            } else {
                RBNode y = leftOf(parentOf(parentOf(rbNode)));//叔叔节点
                if (colorOf(y) == RED) {//如果叔叔节点存在，且为红
                    //变色处理
                    setColor(parentOf(rbNode), BLOCK);
                    setColor(y, BLOCK);
                    setColor(parentOf(y), RED);
                    //递归处理
                    rbNode = parentOf(parentOf(rbNode));
                } else {
                    //如果呈现三角形
                    if (leftOf(parentOf(rbNode)) == rbNode) {
                        //对父节点右旋
                        rightRBNode(parentOf(rbNode));
                    }
                    //对爷爷节点 左旋  变色
                    setColor(parentOf(rbNode), BLOCK);
                    setColor(parentOf(parentOf(rbNode)), RED);
                    leftRBNode(parentOf(parentOf(rbNode)));
                }
            }
        }
        //

        //根节点都为黑色节点
        setColor(root, BLOCK);

    }

    //左旋
    private void leftRBNode(RBNode p) {
        if (p != null) {
            RBNode r = p.right;
            //第一步，让p的右节点指向p的右节点的左节点
            p.right = r.left;
            if (r.left != null) {
                //如果r的左节点不为空，让r的左节点的父节点指向p
                r.left.parent = p;
            }
            //2.r的父节点指向p的父节点，如果p的父节点不为空，就将p的父节点的左或右节点指向r,否者将r设置为root节点
            r.parent = p.parent;
            if (p.parent == null) {
                r = root;
            } else if (p == p.parent.left) {
                p.parent.left = r;
            } else {
                p.parent.right = r;
            }
            //3 将p的父节点指向r ,r的左节点指向p 即设置p为r的左子节点
            p.parent = r;
            r.left = p;
        }
    }

    //右旋
    private void rightRBNode(RBNode p) {
        if (p != null) {
            RBNode l = p.right;
            //第一步，让p的右节点指向p的右节点的左节点
            p.left = l.right;
            if (l.right != null) {
                //如果r的左节点不为空，让r的左节点的父节点指向p
                l.right.parent = p;
            }
            //2.r的父节点指向p的父节点，如果p的父节点不为空，就将p的父节点的左或右节点指向r,否者将r设置为root节点
            l.parent = p.parent;
            if (p.parent == null) {
                l = root;
            } else if (p == p.parent.left) {
                p.parent.left = l;
            } else {
                p.parent.right = l;
            }
            //3 将p的父节点指向l ,l的左节点指向p 即设置p为l的右子节点
            p.parent = l;
            l.right = p;
        }
    }

    public void setRoot(RBNode root) {
        this.root = root;
    }

    static class RBNode<K extends Comparable<K>, V> {
        private RBNode parent;
        private RBNode left;
        private RBNode right;
        private boolean color;
        private K key;
        private V value;

        public RBNode(RBNode parent, RBNode left, RBNode right, boolean color, K key, V value) {
            this.parent = parent;
            this.left = left;
            this.right = right;
            this.color = color;
            this.key = key;
            this.value = value;
        }

        public RBNode(K key, V value, RBNode parent) {
            this.parent = parent;
            this.key = key;
            this.value = value;
        }

        public RBNode() {
        }

        public RBNode getParent() {
            return parent;
        }

        public void setParent(RBNode parent) {
            this.parent = parent;
        }

        public RBNode getLeft() {
            return left;
        }

        public void setLeft(RBNode left) {
            this.left = left;
        }

        public RBNode getRight() {
            return right;
        }

        public void setRight(RBNode right) {
            this.right = right;
        }

        public boolean isColor() {
            return color;
        }

        public void setColor(boolean color) {
            this.color = color;
        }

        public K getKey() {
            return key;
        }

        public void setKey(K key) {
            this.key = key;
        }

        public V getValue() {
            return value;
        }

        public void setValue(V value) {
            this.value = value;
        }
    }
}
